import csv
import os
import warnings
from inspect import currentframe, getframeinfo

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from multiprocessing import Process

import scipy.io as scio

import numpy as np
import torch
from DSPPytorch import *

import util
BASE_DIR = os.path.dirname(__file__)

def LE(lp_range=[1]):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    lp_range = [*lp_range]
    Qcache = np.zeros([len(lp_range)])
    pick_sym_num = 131072
    modOrder = 4
    result_save_dir= os.path.join(BASE_DIR, 'result/{}'.format(experiment_name))
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)

    '''将结果写入CSV文件'''
    with open(os.path.join(result_save_dir, 'Q_results_under_{}.csv'.format(experiment_name)),
        'w') as f:
        csv_writer = csv.writer(f)
        csv_writer. writerow([
            'lp', 'BER', 'Q factor', 'Q^2 factor'
        ])

        '''进入循环，对每个发射功率下的信号进行均衡'''
        for lpIndx, lp in enumerate(lp_range):
            dataPath = os.path.join(
                BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/BPS_4_seed_2_dBm_{lp:d}_Loops_8.mat'.format(lp=lp)
            )
            '''读取文件信息'''
            data = scio.loadmat(dataPath)
            symbolRate = 32e9
            bitsPerSymbol = data['BitsPerSymbol'][0, 0]
            spanNum = data['Loops_N'][0, 0] * 5
            rolloff = 0.01
            lp = data['dBm_N'][0, 0]
            prbsx = data['prbsx']
            prbsy = data['prbsy']
            spanLen = 80e3
            L = spanLen * spanNum
            DL = 17e-6 * L / 2  # 预补偿了50%的色散
            sigx = torch.from_numpy(data['signalx'])
            sigy = torch.from_numpy(data['signaly'])
            sigx = sigx.reshape(1, 1, -1)
            sigy = sigy.reshape(1, 1, -1)
            sig = torch.cat([sigx, sigy], dim=1)

            sigxb2b = torch.from_numpy(data['signalx_label'])
            sigyb2b = torch.from_numpy(data['signaly_label'])
            sigxb2b = sigxb2b.reshape(1, 1, -1)
            sigyb2b = sigyb2b.reshape(1, 1, -1)
            sigb2b = torch.cat([sigxb2b, sigyb2b], dim=1)

            constellations = torch.from_numpy(util.CONST_16QAM)

            '''实例化后端的DSP类'''
            # mf = MatchedFilterLayer(alpha=rolloff, span=100, sample_factor=8, case_num=2)  # span: Correlated symbols
            edc = EDCLayer(symbolRate=symbolRate, DL=DL, case_num=2)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            if torch.cuda.is_available():
                sig = sig.cuda()
                sigb2b = sigb2b.cuda()
                lms = lms.cuda()
                edc = edc.cuda()

            '''DSP procedure for the signal through the link'''
            # sig = sig.cpu()
            # sig = mf(sig)
            # if torch.cuda.is_available():
            #     sig = sig.cuda() # 重新放回cuda上
            sig = sig[..., 1::4]  # resampling
            sig = sig[..., 0:pick_sym_num * 2]
            prbsx = prbsx[..., 0:pick_sym_num * modOrder]
            prbsy = prbsy[..., 0:pick_sym_num * modOrder]
            sig = edc(sig)

            '''后端的TDE均衡PMD以及PDL等，是自适应的方法'''
            sig = sig.cpu()
            sig = FIRLayer.time_vary_infer(sig, err_mode='Godard', iter_num=2, lr=5e-4, tap=13)
            if torch.cuda.is_available():
                sig = sig.cuda()  # 重新放回cuda上
            '''DownSampling and PhaseRecovery'''
            sig = pr(sig[..., 1::2])
            '''DD-LMS'''
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            '''与Prbs进行比较计算误码率'''
            sig = sig.cpu().data.numpy().squeeze()
            prbs = np.concatenate([prbsx, prbsy], axis=0)
            sig, ber, good_rotate = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            ber = np.mean(ber)
            Q = util.ber2q(ber)
            Q2 = util.ber2Q2(ber)
            Qcache[lpIndx] = Q

            print(f'lp {lp} done, Q factor: {Q}')

            # '''绘制星座图'''
            # p = Process(target=util.colorfulPlot, args=(sig[0, np.newaxis, ...], prbsx), kwargs={'show': True})
            # p.start()

            '''保存结果'''
            csv_writer.writerow([lp, ber, Q, Q2])

            '''end'''

def LE_debug(lp_range=[1]):
    lp_range = [*lp_range]
    Qcache = np.zeros([len(lp_range)])
    pick_sym_num = 131072
    modOrder = 4

    '''进入循环，对每个发射功率下的信号进行均衡'''
    for lpIndx, lp in enumerate(lp_range):
        dataPath = os.path.join(
            BASE_DIR, 'data/simulation/Test/BPS_4_seed_1_SymbolRate_32000000000_dBm_2_Loops_5_NumberOfChannels_1_OSNR.mat'
        )
        '''读取文件信息'''
        data = scio.loadmat(dataPath)
        symbolRate = 32e9
        bitsPerSymbol = data['BitsPerSymbol'][0, 0]
        spanNum = data['Loops_N'][0, 0]
        rolloff = 0.01
        lp = data['dBm_N'][0, 0]
        prbsx = data['prbsx']
        prbsy = data['prbsy']
        spanLen = 80e3
        L = spanLen * spanNum
        DL = 17e-6 * L
        sigx = torch.from_numpy(data['signalx'])
        sigy = torch.from_numpy(data['signaly'])
        sigx = sigx.reshape(1, 1, -1)
        sigy = sigy.reshape(1, 1, -1)
        sig = torch.cat([sigx, sigy], dim=1)

        sigxb2b = torch.from_numpy(data['signalx_label'])
        sigyb2b = torch.from_numpy(data['signaly_label'])
        sigxb2b = sigxb2b.reshape(1, 1, -1)
        sigyb2b = sigyb2b.reshape(1, 1, -1)
        sigb2b = torch.cat([sigxb2b, sigyb2b], dim=1)

        constellations = torch.from_numpy(util.CONST_16QAM)

        '''实例化后端的DSP类'''
        # mf = MatchedFilterLayer(alpha=rolloff, span=100, sample_factor=8, case_num=2)  # span: Correlated symbols
        edc = EDCLayer(symbolRate=symbolRate, DL=DL, case_num=2)
        pr = PhaseRecLayer(1024)
        lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

        if torch.cuda.is_available():
            sig = sig.cuda()
            sigb2b = sigb2b.cuda()
            lms = lms.cuda()
            edc = edc.cuda()

        '''DSP procedure for the signal through the link'''
        # sig = sig.cpu()
        # sig = mf(sig)
        # if torch.cuda.is_available():
        #     sig = sig.cuda() # 重新放回cuda上
        sig = sig[..., 1::4]  # resampling
        sig = sig[..., 0:pick_sym_num * 2]
        prbsx = prbsx[..., 0:pick_sym_num * modOrder]
        prbsy = prbsy[..., 0:pick_sym_num * modOrder]
        sig = edc(sig)

        '''后端的TDE均衡PMD以及PDL等，是自适应的方法'''
        sig = sig.cpu()
        sig = FIRLayer.time_vary_infer(sig, err_mode='Godard', iter_num=2, lr=5e-4, tap=13)
        if torch.cuda.is_available():
            sig = sig.cuda()  # 重新放回cuda上
        '''DownSampling and PhaseRecovery'''
        sig = pr(sig[..., 1::2])
        '''DD-LMS'''
        lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                remain=2048, lr=5e-4)
        sig = lms(sig)

        '''与Prbs进行比较计算误码率'''
        sig = sig.cpu().data.numpy().squeeze()
        prbs = np.concatenate([prbsx, prbsy], axis=0)
        sig, ber, good_rotate = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

        ber = np.mean(ber)
        Q = util.ber2q(ber)
        Q2 = util.ber2Q2(ber)
        Qcache[lpIndx] = Q

        print(f'lp {lp} done, Q factor: {Q}')

        '''绘制星座图'''
        p = Process(target=util.colorfulPlot, args=(sig[0, np.newaxis, ...], prbsx), kwargs={'show': True})
        p.start()


if __name__ == '__main__':
    warnings.filterwarnings("ignore")

    # LE(lp_range=[-2, -1, 0, 1, 2, 3, 4])

    '''亲自动手调试自己生成的仿真数据'''
    LE_debug(lp_range=[2])
